# coding=utf-8
# Copyright (C) 2020-2021 PowerLZY.
# This file is part of Bold-Falcon - https://github.com/PowerLZY/Bold-Falcon
# See the file 'docs/LICENSE' for copying permission.

import os
import pkgutil
import importlib
import inspect
import logging
from collections import defaultdict
from distutils.version import StrictVersion

from lib.cuckoo.common.abstracts import Auxiliary, Machinery, LibVirtMachinery, Processing, Detection
from lib.cuckoo.common.abstracts import Report, Signature
from lib.cuckoo.common.config import Config
from lib.cuckoo.common.constants import CUCKOO_ROOT, CUCKOO_VERSION
from lib.cuckoo.common.exceptions import CuckooCriticalError
from lib.cuckoo.common.exceptions import CuckooOperationalError
from lib.cuckoo.common.exceptions import CuckooProcessingError
from lib.cuckoo.common.exceptions import CuckooReportError
from lib.cuckoo.common.exceptions import CuckooDependencyError
from loader import Loader
from lib.cuckoo.core.startup import init_logging, init_modules

log = logging.getLogger(__name__)

_modules = defaultdict(list)

class RunProcessing(object):
    """Analysis Results Processing Engine.

    This class handles the loading and execution of the processing modules.
    It executes the enabled ones sequentially and generates a dictionary which
    is then passed over the reporting engine.
    """

    def __init__(self, task):
        """@param task: task dictionary of the analysis to process."""
        self.task = task
        self.analysis_path = os.path.join(CUCKOO_ROOT, "storage", "analyses", str(task["id"]))
        self.baseline_path = os.path.join(CUCKOO_ROOT, "storage", "baseline")
        self.cfg = Config("processing")

    def process(self, module, results):
        """Run a processing module.
        @param module: processing module to run.
        @param results: results dict.
        @return: results generated by module.
        """
        # Initialize the specified processing module.
        try:
            current = module()
        except:
            log.exception("Failed to load the processing module "
                          "\"{0}\":".format(module))
            return None, None

        # Extract the module name.
        module_name = inspect.getmodule(current).__name__
        if "." in module_name:
            module_name = module_name.rsplit(".", 1)[1]

        try:
            options = self.cfg.get(module_name)
        except CuckooOperationalError:
            log.debug("Processing module %s not found in configuration file",
                      module_name)
            return None, None

        # If the processing module is disabled in the config, skip it.
        if not options.enabled:
            return None, None

        # Give it the path to the baseline directory.
        current.set_baseline(self.baseline_path)
        # Give it the path to the analysis results.
        current.set_path(self.analysis_path)
        # Give it the analysis task object.
        current.set_task(self.task)
        # Give it the options from the relevant processing.conf section.
        current.set_options(options)
        # Give the results that we have obtained so far.
        current.set_results(results)

        try:
            # Run the processing module and retrieve the generated data to be
            # appended to the general results container.
            data = current.run()

            log.debug("Executed processing module \"%s\" on analysis at "
                      "\"%s\"", current.__class__.__name__, self.analysis_path)

            # If succeeded, return they module's key name and the data.
            return current.key, data
        except CuckooDependencyError as e:
            log.warning("The processing module \"%s\" has missing dependencies: %s", current.__class__.__name__, e)
        except CuckooProcessingError as e:
            log.warning("The processing module \"%s\" returned the following "
                        "error: %s", current.__class__.__name__, e)
        except:
            log.exception("Failed to run the processing module \"%s\" for task #%d:",
                          current.__class__.__name__, self.task["id"])

        return None, None

    def run(self):
        """Run all processing modules and all signatures.
        @return: processing results.
        """
        # This is the results container. It's what will be used by all the
        # reporting modules to make it consumable by humans and machines.
        # It will contain all the results generated by every processing
        # module available. Its structure can be observed through the JSON
        # dump in the analysis' reports folder. (If jsondump is enabled.)
        # We friendly call this "fat dict".
        results = {
            "_temp": {},
        }

        # Order modules using the user-defined sequence number.
        # If none is specified for the modules, they are selected in
        # alphabetical order.
        processing_list = list_plugins(group="processing")

        # If no modules are loaded, return an empty dictionary.
        if processing_list:
            processing_list.sort(key=lambda module: module.order)

            # Run every loaded processing module.
            for module in processing_list:
                key, result = self.process(module, results)

                # If the module provided results, append it to the fat dict.
                if key and result:
                    results[key] = result
        else:
            log.info("No processing modules loaded")

        results.pop("_temp", None)

        # Return the fat dict.
        return results

class RunDetection(object):
    """
    描述

    :param task: task dictionary of the analysis to detection.
    :param results: results dict
    """

    def __init__(self, results):

        #self.task = task
        #self.analysis_path = os.path.join(CUCKOO_ROOT, "storage", "analyses", str(task["id"]))
        #self.baseline_path = os.path.join(CUCKOO_ROOT, "storage", "baseline")
        self.cfg = Config("detection")

        self.results = results
        self.predict = {}

    def process(self, module):
        """
        Run a detection module.

        @param module: detection module to run.
        @param results: results dict.
        @return: results generated by module.
        """
        # Initialize the specified detection module.
        try:
            current = module()
        except:
            log.exception("Failed to load the detection module "
                          "\"{0}\":".format(module))
            return None, None

        # Extract the module name.
        module_name = inspect.getmodule(current).__name__
        if "." in module_name:
            module_name = module_name.rsplit(".", 1)[1]

        try:
            options = self.cfg.get(module_name)
        except CuckooOperationalError:
            log.debug("Detection module %s not found in configuration file", module_name)
            return None, None

        # If the Detection module is disabled in the config, skip it.
        if not options.enabled:
            return None, None

        # ---- Detection module 需要的配置 ----
        # Give it the path to the baseline directory.
        # rrent.set_baseline(self.baseline_path)
        # Give it the path to the analysis results.
        # current.set_path(self.analysis_path)
        # Give it the analysis task object.
        # current.set_task(self.task)
        # Give it the options from the relevant Detection.conf section.
        #current.set_options(options)
        # Give the results that we have obtained so far.
        current.load_instance(self.results)
        # ---- Detection module 需要的配置 ----

        try:
            # Run the Detection module and retrieve the generated data to be
            # appended to the general results container.
            data = current.run()

            log.debug("Executed Detection module \"%s\" on analysis at "
                      "\"%s\"", current.__class__.__name__, self.analysis_path)

            # If succeeded, return they module's key name and the data.
            return current.key, data
        except CuckooDependencyError as e:
            log.warning("The Detection module \"%s\" has missing dependencies: %s", current.__class__.__name__, e)
        except CuckooProcessingError as e:
            log.warning("The Detection module \"%s\" returned the following "
                        "error: %s", current.__class__.__name__, e)
        except:
            log.exception("Failed to run the Detection module \"%s\" for task #%d:",
                          current.__class__.__name__, self.task["id"])

        return None, None

    def run(self):
        """Run all detection modules.
        @return: detection results.
        """
        # This is the results container. It's what will be used by all the
        # reporting modules to make it consumable by humans and machines.
        # It will contain all the results generated by every processing
        # module available. Its structure can be observed through the JSON
        # dump in the analysis' reports folder. (If jsondump is enabled.)
        # We friendly call this "fat dict".

        predict = {}

        # Order modules using the user-defined sequence number.
        # If none is specified for the modules, they are selected in
        # alphabetical order.
        detection_list = list_plugins(group="detection")

        # If no modules are loaded, return an empty dictionary.
        if detection_list:
            detection_list.sort(key=lambda module: module.order)

            # Run every loaded detection module.
            for module in detection_list:
                key, predict = self.process(module)

                # If the module provided results, append it to the ’info‘ dict.
                if key and result:
                    self.results['info']['detection'] = {key:predict}
        else:
            log.info("No detection modules loaded")

""" From plugin.py """
def import_package(package):
    prefix = package.__name__ + "."
    for loader, name, ispkg in pkgutil.iter_modules(package.__path__, prefix):
        import_plugin(name)

def import_plugin(name):
    try:
        module = __import__(name, globals(), locals(), ["dummy"], -1)
    except ImportError as e:
        raise CuckooCriticalError("Unable to import plugin "
                                  "\"{0}\": {1}".format(name, e))
    else:
        load_plugins(module)

def load_plugins(module): # 导入插件

    for name, value in inspect.getmembers(module):
        if inspect.isclass(value):
            if issubclass(value, Auxiliary) and value is not Auxiliary:
                register_plugin("auxiliary", value)
            elif issubclass(value, Machinery) and value is not Machinery and value is not LibVirtMachinery:
                register_plugin("machinery", value)
            elif issubclass(value, Processing) and value is not Processing:
                register_plugin("processing", value)
            elif issubclass(value, Detection) and value is not Detection:
                # 只导入detection的子类
                register_plugin("detection", value)
            elif issubclass(value, Report) and value is not Report:
                register_plugin("reporting", value)
            elif issubclass(value, Signature) and value is not Signature:
                register_plugin("signatures", value)
# 全局注册_modules
def register_plugin(group, name):
    global _modules
    group = _modules.setdefault(group, [])
    group.append(name)
# _modules[group]
def list_plugins(group=None):
    if group:
        return _modules[group]
    else:
        return _modules

def init_modules(machinery=True):
    """Initializes plugins."""
    log.debug("Importing modules...")

    # Import all processing modules.
    import modules.processing
    import_package(modules.processing)

    # ToDo:Import all processing modules.
    import modules.detection
    import_package(modules.detection)





if __name__ == "__main__":

    # The first stage is to load the data from the directory holding all the JSONs
    init_modules()
    loader = Loader()
    loader.load_binaries_dir("../sample_data/dict")
    result = loader.binaries['133'].report
    # Then we extract all the relevant information from the loaded samples.
    #features_dict = loader.get_features()

    RunDetection(results=result).run()